% Crank-Nicolson 法求解二维Schrodinger方程
% i \hbar \pdv{u}{t} = -\hbar^2/(2m) \pdv[2]{u}{x}
% \pdv{u}{t} = i\hbar/(2m) \pdv[2]{u}{x}
% \frac{u^{(k+1)_i}-u^{(k)}_i}{\Delta t}= i\hbar/(2m) 1/2(\frac{u^{(k+1)}_{i+1}-2u^{(k+1)}_i+u^{(k)}_{i-1}}{(\Delta x)^2}+\frac{u^{(k)}_{i+1}-2u^{(k)}_i+u^{(k)}_{i-1}}{(\Delta x)^2})
% u^{(k+1)_i}-u^{(k)}_i = \alpha(u^{(k+1)}_{i+1}-2u^{(k+1)}_i+u^{(k)}_{i-1} + u^{(k)}_{i+1}-2u^{(k)}_i+u^{(k)}_{i-1}) (\alpha = i\hbar/(2m) \cdot 1/2 \cdot \Delta t/(Delta x)^2)
% u^{(k+1)_i} - \alpha(u^{(k+1)}_{i+1}-2u^{(k+1)}_i+u^{(k)}_{i-1}) = u^{(k)}_i + \alpha(u^{(k)}_{i+1}-2u^{(k)}_i+u^{(k)}_{i-1})
% A u^{(k+1)} = B u^{(k)}
% 参考：https://wuli.wiki/online/CraNic.html
% 可能有bug
% Gitee Repo

clc
clear

L = 1;
dt = 0.001
dx = 0.02;

[x y] = meshgrid(-L:dx:L);
n = size(x, 1);
u0 = zeros(n,n);
u1 = zeros(n,n);
% 势能场
V = zeros(n,n);

u0 = exp(-500*(x.^2+y.^2));
_u0 = u0(:);

A = spalloc(n^2, n^2, 5*n^2);
B = spalloc(n^2, n^2, 5*n^2);

hbar = 1;
m = 2; %质量
alpha = i*hbar/(2*m)*1/2*dt/(dx)^2;
beta = i/hbar*dt/2;

for i = 1:n
    for j = 1:n
        _k = i+(j-1)*n;

        if i == 1 || i == n || j == 1 || j == n %如果是边界
            A(_k,_k) = 1;
            B(_k,_k) = 1;
        else
            A(_k,_k) = 1+4*alpha+beta*V(i,j);
            A(_k,_k+1) = -alpha;
            A(_k,_k-1) = -alpha;
            A(_k,_k+n) = -alpha;
            A(_k,_k-n) = -alpha;

            B(_k,_k) = 1-4*alpha-beta*V(i,j);
            B(_k,_k+1) = alpha;
            B(_k,_k-1) = alpha;
            B(_k,_k+n) = alpha;
            B(_k,_k-n) = alpha;
        end
    end
end
printf('A is constructed\n')

figure();

for tick = 1:500

    _u1 = A\(B*_u0);
    _u0 = _u1;

    if mod(tick,2)==0
        clf
        hold on
        axis equal
        axis([-L L -L L 0 1])
        caxis([0 1])

        u1 = reshape(_u1,n,n);
        u_squared = real(u1).^2+imag(u1).^2;
        h = surf(x,y,100*u_squared);
        set(h,'edgecolor','none');

        xlabel('x');
        ylabel('y');
        drawnow
        pause(0.01)
        disp(tick)
        disp(sum(u_squared(:)))
    end
end

